package com.shujia.flink.state;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.annotation.JSONField;
import lombok.AllArgsConstructor;
import lombok.Data;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;


public class Demo3ListState {
    public static void main(String[] args) throws Exception{
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        //读取数据
        KafkaSource<String> kafkaSource = KafkaSource.<String>builder()
                //kafka broker列表
                .setBootstrapServers("master:9092,node1:9092,node2:9092")
                //指定topic
                .setTopics("cars")
                //消费者组
                .setGroupId("Demo3ListState")
                //读取数据的位置
                .setStartingOffsets(OffsetsInitializer.latest())
                //数据格式
                .setValueOnlyDeserializer(new SimpleStringSchema())
                .build();

        //使用kafka source
        DataStream<String> lines = env
                .fromSource(kafkaSource, WatermarkStrategy.noWatermarks(), "kafka source");

        //解析数据
        DataStream<Car> cars = lines.map(line -> JSON.parseObject(line, Car.class), TypeInformation.of(Car.class));

        /*
         * 实时统计城市的平均车速
         */

        KeyedStream<Car, String> keyBys = cars.keyBy(Car::getCity);


        DataStream<CityAvgSpeed> avgSpeeds = keyBys.process(new KeyedProcessFunction<String, Car, CityAvgSpeed>() {

            //ListState：为每一个key在状态中保存一个集合
            ListState<Integer> speedsState;

            @Override
            public void open(Configuration parameters) throws Exception {
                RuntimeContext context = getRuntimeContext();
                //创建状态描述对象
                ListStateDescriptor<Integer> listStateDescriptor = new ListStateDescriptor<>("speeds", Types.INT);
                //创建状态
                speedsState = context.getListState(listStateDescriptor);
            }

            @Override
            public void processElement(Car car,
                                       KeyedProcessFunction<String, Car, CityAvgSpeed>.Context ctx,
                                       Collector<CityAvgSpeed> out) throws Exception {
                Integer speed = car.getSpeed();
                //将每个车辆的速度保存在状态中
                speedsState.add(speed);

                //从状态中获取数据
                Iterable<Integer> speeds = speedsState.get();
                //循环计算总的车速和总的车流量
                Double sumSpeed = 0.0;
                Integer count = 0;
                for (Integer spe : speeds) {
                    sumSpeed += spe;
                    count++;
                }
                //计算平均车速
                double avgSpeed = sumSpeed / count;

                //将结果发送到下游
                out.collect(new CityAvgSpeed(car.getCity(), avgSpeed));
            }
        });

        avgSpeeds.print();

        env.execute();


    }

    @Data
    @AllArgsConstructor
    static class Car {
        @JSONField(name = "tollgate_id")
        private String tollgateId;
        @JSONField(name = "license_plate")
        private String licensePlate;
        private String timestamp;
        private Double longitude;
        private Double latitude;
        private Integer speed;
        private String city;
        private String road;
        private String district;
    }

    @Data
    @AllArgsConstructor
    static class CityAvgSpeed {
        private String city;
        private Double avgSpeed;
    }
}
